
# lance_db_utils.py
import lancedb
import pyarrow as pa
from typing import List, Dict, Any, Optional
import pandas as pd

class LanceDBManager:
    """
    LanceDB数据库操作工具类
    """

    def __init__(self, uri: str, api_key: Optional[str] = None, region: Optional[str] = None):
        """
        初始化LanceDB连接
        
        Args:
            uri: 数据库URI
            api_key: API密钥（可选）
            region: 区域（可选）
        """
        if api_key and region:
            self.db = lancedb.connect(uri=uri, api_key=api_key, region=region)
        else:
            self.db = lancedb.connect(uri)

    def create_table(self, table_name: str, data: List[Dict] | pd.DataFrame | pa.Table):
        """
        创建表并插入数据
        
        Args:
            table_name: 表名
            data: 数据，支持列表、DataFrame或Arrow表
            
        Returns:
            Table: LanceDB表对象
        """
        return self.db.create_table(table_name, data=data)

    def open_table(self, table_name: str):
        """
        打开已存在的表
        
        Args:
            table_name: 表名
            
        Returns:
            Table: LanceDB表对象
        """
        return self.db.open_table(table_name)

    def list_tables(self) -> List[str]:
        """
        列出所有表名
        
        Returns:
            List[str]: 表名列表
        """
        return self.db.table_names()

    def drop_table(self, table_name: str):
        """
        删除表
        
        Args:
            table_name: 要删除的表名
        """
        self.db.drop_table(table_name)

    def insert_data(self, table_name: str, data: List[Dict] | pd.DataFrame):
        """
        向表中插入数据
        
        Args:
            table_name: 表名
            data: 要插入的数据
        """
        table = self.open_table(table_name)
        table.add(data)

    def search_vector(self, table_name: str, vector: List[float], limit: int = 10):
        """
        向量搜索
        
        Args:
            table_name: 表名
            vector: 查询向量
            limit: 返回结果数量限制
            
        Returns:
            查询结果
        """
        table = self.open_table(table_name)
        return table.search(vector).limit(limit).to_pandas()

    def filter_query(self, table_name: str, filter_expr: str, limit: int = 10):
        """
        条件查询
        
        Args:
            table_name: 表名
            filter_expr: 过滤表达式
            limit: 返回结果数量限制
            
        Returns:
            查询结果
        """
        table = self.open_table(table_name)
        return table.search().where(filter_expr).limit(limit).to_pandas()

    def get_table_stats(self, table_name: str) -> Dict[str, Any]:
        """
        获取表统计信息
        
        Args:
            table_name: 表名
            
        Returns:
            Dict: 表统计信息
        """
        table = self.open_table(table_name)
        return {
            "name": table_name,
            "schema": str(table.schema),
            "count": table.count_rows()
        }

# 使用示例
if __name__ == "__main__":
    # 初始化数据库连接
    db_manager = LanceDBManager(
        uri="db://default-58ja3w",
        api_key="sk_XO3CIPKXRVET3KANC3RRLDUN2YDR4HANNNSJMY3DB2RVPYN3VDFQ====",
        region="us-east-1"
    )

    # 示例数据
    sample_data = [
        {"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
        {"vector": [5.9, 26.5], "item": "bar", "price": 20.0},
    ]

    # 创建表
    table_name = "my_table"
    db_manager.create_table(table_name, data=sample_data)

    # 插入更多数据
    new_data = [
        {"vector": [1.2, 3.4], "item": "baz", "price": 15.0}
    ]
    db_manager.insert_data(table_name, new_data)

    # 向量搜索
    results = db_manager.search_vector(table_name, [3.1, 4.1], limit=5)
    print("Vector search results:")
    print(results)

    # 条件查询
    filtered_results = db_manager.filter_query(table_name, "price > 10", limit=5)
    print("Filtered results:")
    print(filtered_results)

    # 获取表统计信息
    stats = db_manager.get_table_stats(table_name)
    print("Table stats:")
    print(stats)